from torch.utils.data import Dataset
from PIL import Image
import os
import torchvision.transforms as transforms
import torch
from cnn_net import CNN_NET
import pandas as pd
import warnings
warnings.filterwarnings("ignore")

class MyDataset(Dataset):  #继承Dataset
    def __init__(self, img_path_dir, label_path_dir,transform=None):  #初始化一些属性
        self.img_path_dir = img_path_dir  #文件路径
        self.label_path_dir = label_path_dir
        self.transform = transform  #对图形进行处理，如标准化、截取、转换等
        self.images = os.listdir(self.img_path_dir)  #把路径下的所有文件放在一个列表中
        self.labels = pd.read_csv(self.label_path_dir)

    def __len__(self):  #返回整个数据集的大小
        return len(self.images)

    def __getitem__(self, index):  #根据索引index返回图像及标签
        image_index = self.images[index]  #根据索引获取图像文件名称
        img_path = os.path.join(self.img_path_dir, image_index)  #获取图像的路径或目录
        img = Image.open(img_path).convert('RGB')  # 读取图像

        label_index=image_index.split('.png')[0]
        # print(label_index)
        label = self.labels.loc[self.labels['ID'] == int(label_index)].iloc[0][1:3]

        if self.transform is not None:
            img = self.transform(img)

        return img, torch.Tensor(label)


if __name__ == '__main__':
    myTransforms = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize((20, 50))
    ])

    train_data = MyDataset(r'.\data\train\img',r'.\data\train\label\labels.csv', transform=myTransforms)

    train_data_size = len(train_data)
    print(train_data.__getitem__(23)[0].shape)






